#!/usr/bin/env python
# -*- coding: UTF-8 -*-
'''
@Project ：broadview_coss 
@File    ：consumers.py
@IDE     ：PyCharm 
@Author  ：kdly
@Date    ：2025/2/13 上午11:07 
'''
import json
import threading
import time
import logging
import traceback
from socket import timeout

from threading import Thread

from asgiref.sync import async_to_sync
from channels.generic.websocket import WebsocketConsumer
from django.conf import settings
from guacamole.client import GuacamoleClient
from jwt import PyJWTError

from hosts_mgr.models import HostList
from utils.encrypt import Password

from django.core.cache import cache

from rest_framework_jwt.settings import api_settings

from utils.single_user_login import SingleLogin

logger = logging.getLogger(__name__)


class StreamConsumer:
    def __init__(self, websocket):
        self.ssh_channel = None
        self.websocket = websocket

    def connect(self, host_ip, host_port, sys_user_name, sys_user_passwd, term='xterm', cols=140, rows=50):
        # 延迟导入,解决："ImportError: cannot import name 'Message' from partially initialized module 'paramiko.message' (most
        # likely due to a circular import" 问题
        from paramiko import SSHClient, AutoAddPolicy
        ssh_client = SSHClient()
        ssh_client.set_missing_host_key_policy(AutoAddPolicy())
        try:
            ssh_client.connect(host_ip, host_port, sys_user_name, sys_user_passwd, timeout=10)
            logger.info(f"{sys_user_name}@{host_ip} 连接成功")
        except Exception as e:
            error_msg = {
                "code": 500,
                "type": "connection_error",
                "data": str(e)
            }
            self.websocket.send(json.dumps(error_msg))
            logger.info(f"{sys_user_name}@{host_ip} 连接失败")
            return False

        transport = ssh_client.get_transport()
        self.ssh_channel = transport.open_session()
        self.ssh_channel.get_pty(term=term, width=cols, height=rows)
        self.ssh_channel.invoke_shell()
        msg = f"开始连接到{sys_user_name}@{host_ip} \r\n"
        self.send_to_ws_mes(msg)
        for i in range(2):
            msg = self.ssh_channel.recv(1024).decode('utf-8', 'ignore')
            self.send_to_ws_mes(msg)

    def close(self):
        try:
            self.websocket.close()
            self.ssh_channel.close()
        except Exception as e:
            pass

    def send_to_ws_mes(self, data):
        message = {
            "code": 200,  # 状态码
            "type": "stdout",  # 消息类型
            "data": data  # 实际数据
        }
        self.websocket.send(json.dumps(message))

    def _ws_to_ssh(self, data):
        try:
            self.ssh_channel.send(data)
        except OSError as e:
            self.close()

    def _ssh_to_ws(self):
        try:
            while not self.ssh_channel.exit_status_ready():
                data = self.ssh_channel.recv(1024)  # .decode('utf-8')
                try:
                    decoded_data = data.decode('utf-8', errors='replace')  # 使用错误替换策略
                except UnicodeDecodeError:
                    decoded_data = data.decode('latin-1', errors='replace')  # 备选解码方案
                self.send_to_ws_mes(decoded_data)
        except Exception as e:
            error_msg = {
                "code": 500,
                "type": "system_error",
                "data": str(e)
            }
            self.websocket.send(json.dumps(error_msg))
            self.close()

    def shell(self, data):
        Thread(target=self._ws_to_ssh, args=(data,)).start()
        Thread(target=self._ssh_to_ws).start()

    def resize_pty(self, cols, rows):
        try:
            cols = int(float(cols))  # 兼容数字字符串
            rows = int(float(rows))
            self.ssh_channel.resize_pty(width=max(1, cols), height=max(1, rows))
        except (TypeError, ValueError) as e:
            logger.error(f"Invalid resize parameters: {cols}x{rows}, error: {e}")


class SSHConsumer(WebsocketConsumer):
    def __init__(self, *args, **kwargs):
        super().__init__(args, kwargs)
        self.group_name = None
        self.ssh = None

    def auth(self, token):
        try:
            jwt_decode_handler = api_settings.JWT_DECODE_HANDLER
            res = jwt_decode_handler(token)
            username = res.get('username')
            # ================单用户登录逻辑=====================
            single_login = SingleLogin(username)
            if single_login.get_black_list(token):
                logger.warning(f"用户 {username} 被强制下线 (Token: {token})")
                self.send_to_ws_mes(401, "auth_error", '已有用户登录，当前强制下线')
                return False
            # ================单用户登录逻辑=====================
        except PyJWTError as e:
            logger.error(f"Token认证失败,{e}")
            self.send_to_ws_mes(401, "auth_error", "认证失败")
            return False
        return True

    def send_message(self, event):
        """接收 Channel Layer 消息并发送给 WebSocket 客户端"""
        message = event["message"]
        try:
            self.send(text_data=json.dumps(message))
        except Exception as e:
            logger.error(e)

    def send_to_ws_mes(self, code, msg_type, data):
        try:
            message = {
                "code": code,  # 状态码
                "type": msg_type,  # 消息类型
                "data": data  # 实际数据
            }
            self.send(json.dumps(message))
        except AttributeError as e:
            logger.error(f"WebSocket未正确初始化，无法发送消息: {e}")
        except TypeError as e:
            logger.error(f"消息序列化失败 (类型错误): {e}\n原始数据: {data}")
        except Exception as e:
            logger.error(f"发送消息时发生未知错误: {e}\n{traceback.format_exc()}")

    def connect(self):
        self.accept()
        try:
            token = self.scope["cookies"]["token"]
        except KeyError:
            self.send_to_ws_mes(401, "auth_error", '请重新登陆')
            self.close()
            return
        if not self.auth(token):
            self.send_to_ws_mes(401, "auth_error", '请重新登陆')
            self.close()
            return
        try:
            self.group_name = self.scope['url_route']['kwargs']['group_name']
            # 将当前频道加入频道组
            async_to_sync(self.channel_layer.group_add)(
                self.group_name,
                self.channel_name
            )
            # 加入组后设置活动时间，方便在ssh_client.py中获取self.group_name组是否断开ws连接
            cache.set(f'group_activity:{self.group_name}', time.time(), timeout=1800)
            if self.group_name == "terminal":
                host_info = self.scope["query_string"].decode()
                host_id = int(host_info.split('&')[0].split('=')[1])

                host_obj = HostList.objects.get(id=host_id)
                host_ip = host_obj.connect_ip
                host_port = host_obj.port
                user_name = host_obj.user
                user_pass = Password().decrypt(host_obj.password)

                self.ssh = StreamConsumer(websocket=self)
                self.ssh.connect(host_ip, host_port, user_name, user_pass)
        except Exception as e:
            logger.error(f"WebSocket连接初始化失败: {e}\n{traceback.format_exc()}")
            self.close()

    def close_connection(self, event=''):
        """收到关闭指令时，关闭 WebSocket 连接"""
        self.close()  # 关闭连接

    def disconnect(self, close_code):
        try:
            async_to_sync(self.channel_layer.group_discard)(
                self.group_name,
                self.channel_name
            )
            if self.ssh:
                self.ssh.close()
            # 删除组活动时间记录，如果 self.group_name 组断开了连接，就删除缓存 group_activity:{self.group_name} 键
            cache.delete(f'group_activity:{self.group_name}')
            logger.warning(f'已删除 group_activity:{self.group_name} ws组')
        except TypeError as e:
            logger.error(f"断开连接时发生错误: {e}\n{traceback.format_exc()}")

    def receive(self, text_data=None, bytes_data=None):
        try:
            text_data = json.loads(text_data)
            if text_data.get('flag') == 'resize':
                # 转换cols和rows为整数
                cols = int(text_data['cols'])
                rows = int(text_data['rows'])
                if self.ssh:
                    self.ssh.resize_pty(cols=cols, rows=rows)
            elif text_data.get('type') == 'heartbeat':
                message = {
                    "code": 200,
                    "type": "heartbeat",
                    "data": "pong",
                }
                self.send(json.dumps(message))
            else:
                data = text_data.get('entered_key', '')
                self.ssh.shell(data=data)
        except json.JSONDecodeError as e:
            logger.error(f"接收到的消息格式错误: {text_data}\n错误详情: {e}")
        except KeyError as e:
            logger.error(f"消息缺少必要字段: {e}\n原始数据: {text_data}")
        except Exception as e:
            logger.error(f"处理消息时发生未知错误: {e}\n{traceback.format_exc()}")


class GuacamoleThread(threading.Thread):
    """创建guacamole线程"""

    def __init__(self, message):
        super(GuacamoleThread, self).__init__()
        self._stop_event = threading.Event()
        self.message = message
        self.read_lock = threading.RLock()
        self.write_lock = threading.RLock()

    def stop(self):
        self._stop_event.set()

    def run(self):
        with self.read_lock:
            while not self._stop_event.is_set():
                time.sleep(0.001)
                try:
                    instruction = self.message.client.receive()
                    if instruction:
                        # 检查连接是否已关闭
                        if not self.message.guacamole_thread.is_alive():
                            break
                        self.message.send(instruction)
                except timeout:
                    logging.getLogger().error('连接超时！')
                    break
                except RuntimeError as e:
                    logger.warning(f"连接已关闭，忽略发送指令: {e}")
                    break
                except OSError:
                    break

            # 安全发送关闭指令
            try:
                if not self._stop_event.is_set():
                    self.message.send('0.;')
            except Exception as e:
                logger.warning(f"发送关闭指令失败: {e}")


class GuacamoleConsumer(WebsocketConsumer):
    """创建guacamole连接"""

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.group_name = None
        self.client = GuacamoleClient(settings.GUACD_HOST, settings.GUACD_PORT, timeout=5)
        self.guacamole_thread = GuacamoleThread(self)
        self.__encryption_pass = Password()

    def connect(self):
        self.accept('guacamole')
        token = self.scope["cookies"]["token"]
        if not SSHConsumer().auth(token):
            return
        self.group_name = self.scope['url_route']['kwargs']['group_name']

        server = HostList.objects.get(id=self.scope['url_route']['kwargs']['host_id'])
        print(server.hostname, '=============hostname')
        print(server.rdp_port, '=============rdp_port')
        print(server.user, '=============user')
        print(server.connect_ip, '=============connect_ip')
        query_list = self.scope['query_string'].decode('utf8').split(',')
        width = query_list[0].strip()
        height = query_list[1].strip()
        dpi = query_list[2].strip()
        print(dpi, '=============dpi')
        self.client.handshake(
            protocol='rdp',
            security='any',  # Use NLA security (more secure and commonly supported)
            ignore_cert='true',  # Properly formatted as string
            hostname=server.connect_ip,
            port=server.rdp_port,
            password=self.__encryption_pass.decrypt(server.password),
            username=server.user,
            width=width,
            height=height,
            dpi=dpi,
            clipboard='enable'
        )
        self.send(f'0.,{len(self.group_name)}.{self.group_name};')
        self.guacamole_thread.daemon = True
        self.guacamole_thread.start()

    def disconnect(self, event):
        # 先停止线程
        if hasattr(self, 'guacamole_thread'):
            self.guacamole_thread.stop()
            self.guacamole_thread.join()  # 等待线程完全退出
        # 再关闭连接
        self.client.close()
        super().disconnect(event)

    def receive(self, text_data=None, bytes_data=None):
        with self.guacamole_thread.write_lock:
            self.client.send(text_data)
